"""词向量相关类"""
import numpy as np
from tqdm import tqdm
import pickle

class Word2Vec:
    def __init__(self):
        self.word2id,self.id2word,self.word2vec = self.load_word2vec('D://giga.100.txt')

    def get_word2vec(self,word):
        """
        根据单词获取词向量
        :param word:
        :return:
        如果单词存在词向量，则返回它的词向量，否则返回-1
        """
        if word in self.word2id.keys():
            idx = self.word2id[word]
            return self.word2vec[idx]
        else:
            return -1

    def get_similarity(self,word1,word2):
        """
        计算两个单词的余弦相似度
        :param word1:
        :param word2:
        :return:
        """
        word2vec_1 = self.get_word2vec(word1)
        word2vec_2 = self.get_word2vec(word2)
        if word1 not in self.word2id.keys() or word2 not in self.word2id.keys():
            return -1
        return word2vec_1.dot(word2vec_2)/(np.linalg.norm(word2vec_1)*np.linalg.norm(word2vec_2))

    def load_word2vec(self,filename):
        """
        读取词向量
        :param filename: 词向量文件名
        :return:
        word2id:词映射至索引
        id2word:索引映射至词
        word2vec:词向量矩阵
        """
        #filename = 'D://giga.100.txt'
        word2id = {}
        id2word = {}
        word2vec = []
        with open(filename,'rb') as f:
            count = 0
            for line in tqdm(f):
                tmp = line.decode('utf-8').rstrip('\n').split()
                word2vec.append([float(i) for i in tmp[1:]])
                word2id[tmp[0]] = count
                id2word[count] = tmp[0]
                count+=1
        word2vec = np.array(word2vec)
        return word2id,id2word,word2vec

    def save(self):
        """
        保存
        :return:
        """
        with open("model/WordEmbedding.txt", "wb") as f:
            pickle.dump(self,f)

